from my_bonito.api.api_config import GPT4_CONFIG
from my_bonito.utils.post_processing import extract_sentences, get_stripped_sentences, find_word_in_string
import string
from openai import OpenAI

class GPT4:
    def __init__(self, config: GPT4_CONFIG):
        self.config = config

    def make_requests(self, prompt):
        # print(f"Calling GPT-4 API with config: {self.config.__dict__}")
        client = OpenAI(
            api_key=self.config.api_key,
        )
        response = client.chat.completions.create(
            model=self.config.engine,
            messages=[
                {"role": "system", "content": self.config.system_prompt},
                {"role": "user", "content": prompt}
            ],
            max_tokens=self.config.max_tokens,
            temperature=self.config.temperature,
            top_p=self.config.top_p
        )
        return response

    def post_processing(self, response):
        if response is None or response.choices[0].finish_reason=="length":
            return []
        raw_instructions = get_stripped_sentences(extract_sentences(response.choices[0].message.content))
        instructions = []
        for inst in raw_instructions:
            if inst == "":
                continue
            # filter out too short or too long instructions
            if len(inst) <= 3 or len(inst) > 150:
                continue
            # filter based on keywords that are not suitable for language models.
            if any(find_word_in_string(word, inst) for word in
                   ["image", "images", "graph", "graphs", "picture", "pictures", "file", "files", "map", "maps", "draw",
                    "plot", "go to"]):
                continue
            # We found that the model tends to add "write a program" to some existing instructions, which lead to a lot of such instructions.
            # And it's a bit confusing whether the model need to write a program or directly output the result.
            # Here we filter them out.
            # Note this is not a comprehensive filtering for all programming instructions.
            if inst.startswith("Write a program"):
                continue
            # filter those starting with punctuation
            if inst[0] in string.punctuation:
                continue
            # filter those starting with numbers
            if inst[0].isnumeric():
                continue
            instructions.append(inst)

        return instructions, response

    def get_raw_output(self, prompt):
        return self.make_requests(prompt).choices[0].message.content

    def get_instructions(self, prompt):
        return self.post_processing(self.make_requests(prompt))


